#!/bin/bash

# 机器学习预测服务部署脚本
# 使用方法: ./deploy.sh <服务器IP> <SSH端口>
# 功能: 自动部署机器学习预测API到远程服务器

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 检查参数
if [ $# -ne 2 ]; then
    echo -e "${RED}使用方法: $0 <服务器IP> <SSH端口>${NC}"
    echo -e "${YELLOW}示例: $0 112.86.254.98 3022${NC}"
    exit 1
fi

SERVER_IP=$1
SSH_PORT=$2
REMOTE_USER="root"
PROJECT_NAME="ml-predict-api"
REMOTE_DIR="/opt/$PROJECT_NAME"

echo -e "${BLUE}开始部署机器学习预测服务到服务器 $SERVER_IP:$SSH_PORT${NC}"

# 检查本地Docker是否运行
# if ! docker info > /dev/null 2>&1; then
#     echo -e "${RED}错误: 本地Docker未运行，请先启动Docker${NC}"
#     exit 1
# fi

# 检查SSH连接
echo -e "${YELLOW}检查SSH连接...${NC}"
if ! ssh -p $SSH_PORT -o ConnectTimeout=10 $REMOTE_USER@$SERVER_IP exit 2>/dev/null; then
    echo -e "${RED}错误: 无法连接到服务器 $SERVER_IP:$SSH_PORT${NC}"
    echo -e "${YELLOW}请确保:${NC}"
    echo "1. 服务器IP和端口正确"
    echo "2. SSH密钥已配置或密码认证可用"
    echo "3. 服务器防火墙允许SSH连接"
    exit 1
fi

echo -e "${GREEN}SSH连接成功${NC}"

# 在远程服务器上安装Docker（如果未安装）
echo -e "${YELLOW}检查远程服务器Docker环境...${NC}"
ssh -p $SSH_PORT -t $REMOTE_USER@$SERVER_IP << 'EOF'
if ! command -v docker &> /dev/null; then
    echo "安装Docker..."
    curl -fsSL https://get.docker.com -o get-docker.sh
    sudo sh get-docker.sh
    sudo systemctl start docker
    sudo systemctl enable docker
    rm get-docker.sh
    echo "Docker安装完成"
else
    echo "Docker已安装"
fi

# 确保Docker服务正在运行
echo "检查Docker服务状态..."
if ! sudo systemctl is-active --quiet docker; then
    echo "启动Docker服务..."
    sudo systemctl start docker
    sudo systemctl enable docker
fi

# 等待Docker服务完全启动
echo "等待Docker服务启动..."
sleep 5

# 验证Docker是否正常工作
if sudo docker info > /dev/null 2>&1; then
    echo "Docker服务运行正常"
else
    echo "Docker服务启动失败，尝试重启..."
    sudo systemctl restart docker
    sleep 5
    if sudo docker info > /dev/null 2>&1; then
        echo "Docker服务重启成功"
    else
        echo "Docker服务启动失败，请检查日志"
        sudo systemctl status docker
        exit 1
    fi
fi
EOF

# 安装Docker Compose
echo -e "${YELLOW}检查Docker Compose...${NC}"
ssh -p $SSH_PORT $REMOTE_USER@$SERVER_IP << 'EOF'
# 检查Docker Compose是否已安装
echo "检查Docker Compose..."
if command -v docker-compose &> /dev/null; then
    echo "Docker Compose已安装，版本: $(docker-compose --version)"
else
    echo "Docker Compose未安装，但跳过安装（使用纯Docker命令）"
fi
EOF

# 创建远程目录
echo -e "${YELLOW}创建远程项目目录...${NC}"
ssh -p $SSH_PORT -t $REMOTE_USER@$SERVER_IP "sudo mkdir -p $REMOTE_DIR && sudo chown $REMOTE_USER:$REMOTE_USER $REMOTE_DIR"

# 复制项目文件到远程服务器
echo -e "${YELLOW}上传项目文件...${NC}"
scp -P $SSH_PORT -r . $REMOTE_USER@$SERVER_IP:$REMOTE_DIR/

# 验证模型文件是否正确上传
echo -e "${YELLOW}验证模型文件...${NC}"
ssh -p $SSH_PORT $REMOTE_USER@$SERVER_IP "ls -la $REMOTE_DIR/save_model/"

# 在远程服务器上构建和启动服务
echo -e "${YELLOW}在远程服务器上构建和启动服务...${NC}"
ssh -p $SSH_PORT -t $REMOTE_USER@$SERVER_IP << EOF
cd $REMOTE_DIR

# 停止现有容器（如果存在）
echo "停止现有容器..."
sudo docker stop ml-api-python 2>/dev/null || true
sudo docker rm ml-api-python 2>/dev/null || true

# 检查镜像是否存在
echo "检查Docker镜像..."
if sudo docker images | grep -q "ml-predict-api.*python"; then
    echo "镜像已存在，跳过构建"
else
    echo "镜像不存在，开始构建..."
    sudo docker build -f Dockerfile -t ml-predict-api:python .
    echo "镜像构建完成"
fi

# 启动容器
echo "启动Docker容器..."
sudo docker run -d -p 5090:5090 --name ml-api-python \
    -v $REMOTE_DIR/save_model:/app/save_model:ro \
    ml-predict-api:python

# 等待服务启动
echo "等待服务启动..."
sleep 10

# 检查服务状态
if sudo docker ps | grep -q "ml-api-python"; then
    echo -e "\${GREEN}服务启动成功！\${NC}"
    echo -e "\${BLUE}服务信息:\${NC}"
    echo "服务地址: http://$SERVER_IP:5090"
    echo "预测接口: http://$SERVER_IP:5090/predict"
    echo ""
    echo -e "\${YELLOW}测试命令:\${NC}"
    echo "curl -X POST http://$SERVER_IP:5090/predict \\"
    echo "  -H 'Content-Type: application/json' \\"
    echo "  -d '{\"feature_0\": 1.2, \"feature_1\": -0.5, \"feature_2\": 0.8, \"feature_3\": 1.1, \"feature_4\": -0.3}'"
    echo ""
    echo -e "\${YELLOW}查看日志:\${NC}"
    echo "sudo docker logs ml-api-python"
    echo ""
    echo -e "\${YELLOW}停止服务:\${NC}"
    echo "sudo docker stop ml-api-python"
else
    echo -e "\${RED}服务启动失败，请检查日志:\${NC}"
    sudo docker logs ml-api-python
    exit 1
fi
EOF

echo -e "${GREEN}部署完成！${NC}"
echo -e "${BLUE}服务已成功部署到 http://$SERVER_IP:5090${NC}"
